import numpy as np
import cv2
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
import pickle

# 加载图像和分区图像
def load_image(image_path, segmentation_path):
    try:
        # 以灰度模式读取图像和分区图像
        image = cv2.imread(image_path, 0)
        segmentation = cv2.imread(segmentation_path, 0)
        
        # 检查图像和分区图像是否读取成功
        if image is None or segmentation is None:
            raise FileNotFoundError("图像或分区图像未找到")
        return image, segmentation
    except Exception as e:
        print(f"加载图像时出错: {e}")
        raise

# 预处理图像和分区数据
def preprocess_data(image, segmentation):
    # 将图像和分区数据展平为一维数组
    X = image.reshape(-1, 1)
    y = segmentation.reshape(-1)
    print("完成从砂岩截面图1及其对应分区中获取X和y")
    return X, y

# 训练随机森林模型
def train_model(X_train, y_train):
    # 初始化随机森林分类器
    clf = RandomForestClassifier(n_jobs=-1, random_state=42)  # 使用多核处理提高效率
    clf.fit(X_train, y_train)
    print("完成train_test_split")
    print("完成随机森林模型clf的训练")    
    return clf

# 保存模型到硬盘
def save_model(clf, model_path):
    # 使用pickle保存模型
    with open(model_path, "wb") as f:
        pickle.dump(clf, f)
    print("模型已成功保存到硬盘")

# 主函数
def main():
    image_path = r"D:/work/tx/Sandstone_1.tif"
    segmentation_path = r"D:/work/tx/Sandstone_1_segment.tif"
    model_path = r"D:/work/tx/clf.pkl"

    # 加载图像和分区数据
    image, segmentation = load_image(image_path, segmentation_path)
    print(f"图像形状: {image.shape}")
    
    # 预处理数据
    X, y = preprocess_data(image, segmentation)

    # 拆分数据集为训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # 训练模型
    clf = train_model(X_train, y_train)

    # 评估模型在测试集上的准确性
    accuracy = clf.score(X_test, y_test)
    print(f"模型在测试集上的准确率: {accuracy:.4f}")

    # 保存训练好的模型
    save_model(clf, model_path)

# 运行主函数
if __name__ == "__main__":
    main()
